''' 3d morphable model example
3dmm parameters --> mesh 
fitting: 2d image + 3dmm -> 3d face
'''
import os, sys
import subprocess
import numpy as np
import scipy.io as sio
from skimage import io
from time import time
import matplotlib.pyplot as plt
import re
import math

sys.path.append('..')
import face3d
from face3d import mesh
from face3d.morphable_model import MorphabelModel
from face3d import mesh
from face3d.morphable_model.fit_multi import MultiFit

import cv2
import dlib
import string

detector = dlib.get_frontal_face_detector()
predictor = dlib.shape_predictor('../Data/shape_predictor_68_face_landmarks.dat')

class my3DFace(object):
    def __init__(self):
        super( my3DFace, self).__init__()
        # --- 1. load model
        self.bfm = MorphabelModel('../Data/BFM.mat')
        print('init bfm model success')

    # 将人脸模型保存到文件中，供C++使用
    def saveModelforC(self, fname):
        sm = self.bfm.model['shapeMU']
        sp = self.bfm.model['shapePC']
        ep = self.bfm.model['expPC']
        fc = self.bfm.model['tri']
        #fc = np.vstack((bfm.model['tri'], bfm.model['tri_mouth']))

        with open(fname, 'wt') as f:
            f.write("3DModel by lxj-8d\n")
            f.write("MP %6d\n" %(sm.shape[0]))
            for i in range(sm.shape[0]):
                f.write('M %8.3f\n'%(sm[i][0]))

            f.write("SP  %6d  %6d\n" %(sp.shape[0], sp.shape[1]))
            for i in range(sp.shape[0]):
                f.write('S')
                for j in range(sp.shape[1]):
                    f.write(' %8.3f' %(sp[i][j]*100000.0))
                f.write('\n')

            f.write("EP  %6d  %6d\n" %(ep.shape[0], ep.shape[1]))
            for i in range(ep.shape[0]):
                f.write('E')
                for j in range(ep.shape[1]):
                    f.write(' %8.3f' %(ep[i][j]))
                f.write('\n')

            f.write("FC  %6d  %6d\n" %(fc.shape[0], fc.shape[1]))
            for i in range(fc.shape[0]):
                f.write('F')
                for j in range(fc.shape[1]):
                    f.write(' %6d' %(fc[i][j]))
                f.write('\n')

            f.write("END\n")

    # 将人脸模型保存到文件中，供webGL使用
    def saveModelforWebGL(self, fname):
        sm = self.bfm.model['shapeMU']
        sp = self.bfm.model['shapePC']
        ep = self.bfm.model['expPC']
        fc = self.bfm.model['tri']
        #fc = np.vstack((bfm.model['tri'], bfm.model['tri_mouth']))

        tp = 1 #bfm.get_tex_para('random')
        colors = self.bfm.generate_colors(tp)
        colors = np.minimum(np.maximum(colors, 0), 1)
        vp = sm.reshape((sm.shape[0]//3, 3))

        with open(fname, 'wt') as f:
            f.write("#3DModel by lxj-8d\n")
            f.write("#MP %6d\n" %(vp.shape[0]))
            for i in range(vp.shape[0]):
                f.write('v %4.3f %4.3f %4.3f %4.3f %4.3f %4.3f\n'%(vp[i][0], vp[i][1], vp[i][2], colors[i][0], colors[i][1], colors[i][2]))

            f.write("#FC  %6d  %6d\n" %(fc.shape[0], fc.shape[1]))
            for i in range(fc.shape[0]):
                f.write('f')
                for j in range(fc.shape[1]):
                    f.write(' %d' %(fc[i][j]+1))
                f.write('\n')

            f.write("#END\n")

    #计算图像的68个特征点， 返回特征点表，为保证解算时的均衡，进行了比例和平移变换，变换参数同时返回
    def calFeatures68(self, name):
        frame = cv2.imread(name)
        frame_new = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)
        # 检测脸部
        dets = detector(frame_new, 1)
        #print("Number of faces detected: {}".format(len(dets)))
        # 查找脸部位置
        # 绘制脸部位置
        x = np.zeros((68,2), dtype=np.float32)
        for i, face in enumerate(dets):
            #cv2.rectangle(frame, (face.left(), face.top()), (face.right(), face.bottom()), (0, 255, 0), 1)
            x0 = (face.left() + face.right())/2
            y0 = (face.top() + face.bottom())/2
            dx = face.left() - face.right()
            dy = face.top() - face.bottom()
            radius = math.sqrt(dx*dx + dy*dy)
            scale = 100.0/radius
            #x0 = 200
            #y0 = 200
            shape = predictor(frame_new, face)
            for j in range(68):
                #xx = int(x[j][0]*2)
                #yy = -int(x[j][1]*2)
                xx = shape.part(j).x - x0
                yy = shape.part(j).y - y0
                x[j][0] = shape.part(j).x - x0
                x[j][1] = y0 - shape.part(j).y
                cv2.circle(frame,(int(xx+x0), int(yy+y0)),1,(0,0,255),2)
                #cv2.putText(frame,str(j),(int(xx+x0), int(yy+y0)),cv2.FONT_HERSHEY_COMPLEX,0.5,(255,0,0),1)
                x[j][0] *= scale
                x[j][1] *= scale
        #cv2.imshow(name,frame)
        #key = cv2.waitKey(1)
        return x, x0, y0, scale
        #xlist = xlist + [x.copy()]

    # 显示识别点与三维点的误差图，测试用
    def dispToler(self, name, res):
        frame = cv2.imread(name)
        frame_new = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)
        # 检测脸部
        dets = detector(frame_new, 1)
        #print("Number of faces detected: {}".format(len(dets)))
        # 查找脸部位置
        # 绘制脸部位置
        x = np.zeros((68,2), dtype=np.float32)
        for i, face in enumerate(dets):
            #cv2.rectangle(frame, (face.left(), face.top()), (face.right(), face.bottom()), (0, 255, 0), 1)
            x0 = (face.left() + face.right())/2
            y0 = (face.top() + face.bottom())/2
            dx = face.left() - face.right()
            dy = face.top() - face.bottom()
            radius = math.sqrt(dx*dx + dy*dy)
            scale = 100.0/radius
            #x0 = 200
            #y0 = 200
            shape = predictor(frame_new, face)
            for j in range(68):
                #xx = int(x[j][0]*2)
                #yy = -int(x[j][1]*2)
                xx = shape.part(j).x - x0
                yy = shape.part(j).y - y0
                x[j][0] = shape.part(j).x - x0
                x[j][1] = y0 - shape.part(j).y
                cv2.circle(frame,(int(xx+x0), int(yy+y0)),1,(0,0,255),2)
                #cv2.putText(frame,str(j),(int(xx+x0), int(yy+y0)),cv2.FONT_HERSHEY_COMPLEX,0.5,(255,0,0),1)
            #disp res
            res /= scale
            for j in range(68):
                xx = res[j][0] + x0
                yy = y0 - res[j][1]
                cv2.circle(frame,(int(xx), int(yy)),1,(255,0,0),2)
        cv2.imshow(name,frame)
        #key = cv2.waitKey(1)

    # 将多张原始图像的颜色合成到一起
    def mergeColor(self, imgname, x0, y0, fitted_sp0, ep, slist, anglist, tlist, scale):
        for index, fname in enumerate(imgname):
            fitted_vertices = self.bfm.generate_vertices(fitted_sp0, ep[index])
            transformed_vertices = self.bfm.transform(fitted_vertices, slist[index], anglist[index], tlist[index])

            #求每个点的法向量
            normal = np.zeros((transformed_vertices.shape[0],3), dtype=np.float32)
            for i, tri in enumerate(self.bfm.triangles):
                p0 = tri[0]
                p1 = tri[1]
                p2 = tri[2]
                a = transformed_vertices[p1] - transformed_vertices[p0]
                b = transformed_vertices[p2] - transformed_vertices[p0]

                v3 = np.array([a[1]*b[2]-a[2]*b[1],a[2]*b[0]-a[0]*b[2],a[0]*b[1]-a[1]*b[0]])
                delta = v3.dot(v3)
                delta = math.sqrt(delta)
                v3 /= delta
                normal[p0] += v3
                normal[p1] += v3
                normal[p2] += v3
            for i in range(normal.shape[0]):
                if normal[i][2] < 0:
                    delta = normal[i].dot(normal[i])
                    delta = math.sqrt(delta)
                    normal[i] /= delta
                    normal[i][2] *= normal[i][2]
                else:
                    normal[i][2] = 0.000001

            transformed_vertices = transformed_vertices[:,:2]
            transformed_vertices /= scale[index]
            if (index == 0):
                colors = np.zeros((transformed_vertices.shape[0], 3), dtype = np.float32)
                wt = np.zeros((transformed_vertices.shape[0],), dtype = np.float32)
            frame = cv2.imread(fname)
            for i in range(transformed_vertices.shape[0]):
                xx = transformed_vertices[i][0]
                yy = transformed_vertices[i][1]
                xx += x0[index]
                yy = y0[index] - yy
                xi = int(xx + 0.5)
                yi = int(yy + 0.5)
                c = frame[yi][xi]/256.0
                colors[i] += c*normal[i][2]
                wt[i] += normal[i][2]
        for i in range(colors.shape[0]):
            if wt[i] > 0.0000000001:
                colors[i] /= wt[i]
        return colors

    #保存三维模型
    def saveFace(self, facename, fitted_sp0, fitted_ep0, colors):
        with open(facename, 'w') as f:
            f.write("3DFace by lxj-8d\n")
            f.write("SP %6d\n" %(fitted_sp0.shape[0]))
            for i in range(fitted_sp0.shape[0]):
                f.write('%8.3f\n'%(fitted_sp0[i]/100000.0))
            f.write("EP  %6d\n" %(fitted_ep0.shape[0]))
            for i in range(fitted_ep0.shape[0]):
                f.write('%8.3f\n' %(fitted_ep0[i]))
            f.write("CL  %8d\n" %(colors.shape[0]))
            for i in range(colors.shape[0]):
                f.write('C %5.3f  %5.3f  %5.3f\n' %(colors[i][0], colors[i][1], colors[i][2]))
                #colors[i] = c
            f.write("END\n")

    #若干张2D图像生成一个三维模型的主函数
    def images_3D(self, pname, resfile):
        #pname = ("e:/8d/data/ab1.jpg", "e:/8d/data/ab3.jpg")
        imgnum = len(pname)
        x0 = np.zeros((imgnum, ), dtype = np.float32)
        y0 = np.zeros((imgnum, ), dtype = np.float32)
        sc0 = np.zeros((imgnum, ), dtype = np.float32)
        xlist = []
        for index,name in enumerate(pname):
            x, x0[index], y0[index], sc0[index] = self.calFeatures68(name)
            xlist = xlist + [x]

        X_ind = self.bfm.kpt_ind
        fits = MultiFit()
        k = 1;

        # fit
        #fitted_sp0, fitted_ep0, fitted_s0, fitted_angles0, fitted_t0 = bfm.fit(x, X_ind, max_iter = 3)
        fitted_sp0, ep, slist, anglist, tlist = fits.fit_points(xlist, X_ind, self.bfm, 3)
        #fitted_angles0 = mesh.transform.matrix2angle(R[k])
        fitted_ep0 = ep[k]
        s = slist[k]
        angles = anglist[k]
        t = tlist[k]

        ##以下进行精度比较，测试用
        #for index, x in enumerate(xlist):
        #    err, res = fits.fit_comparebypost(x, X_ind, bfm, ep[index], slist[index], anglist[index], tlist[index], 3)
        #    print('erroor is %6.3f for file %s'%(err, pname[index]))

        #testimg = [pname[0]]
        #for index, name in enumerate(testimg):
        #    x, xx, yy, sc = calFeatures68(name)
        #    err, res = fits.fit_compare(x, X_ind, bfm, 3)
        #    print('erroor is %6.3f for file %s'%(err, name))
        #    dispToler(name, res)

        colors = self.mergeColor(pname, x0, y0, fitted_sp0, ep, slist, anglist, tlist, sc0)
        self.saveFace(resfile, fitted_sp0, fitted_ep0, colors)

face = my3DFace()
face.saveModelforWebGL("../Data/mod3dweb.obj")

st = 44
path = "../Data/pointing/"
#for ps in ["personne012", "personne021", "personne031", "personne041", "personne051"]:
for ps in ["personne071"]:
    tt = []
    for vi in ["+0"]:
        for index, hi in enumerate(["-30", "-15", "+0", "+15", "+30"]):
            name = path + "vertical"+vi+"/horizon"+hi+"/"+ps+"%d"%(st+index)+vi+hi+".jpg"
            tt = tt + [name]
    pname = tt
    resfile = "%s.txt"%(ps)
    face.images_3D(pname, resfile)

